# Testing decoders.
# 23-10-30

import os
import sys
import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk
import mindspore as ms
import mindspore.nn as ms_nn

from jax import Array, jit
from typing import Optional, Union, List
from jax.nn.initializers import lecun_normal, normal, zeros, ones
from mindspore import load_param_into_net, Tensor, nn, context

sys.path.append(os.path.dirname(sys.path[0]))

# setting context & hyper parameters
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")

from cybertron.readout.decoder import HalveDecoder, ResDecoder
from ms_cybertron.decoder import HalveDecoder as ms_HalveDecoder
from ms_cybertron.decoder import ResidualOutputBlock as ms_ResDecoder

EPSILON = 1e-3
rng = jax.random.PRNGKey(42)

# # test halve decoder
# print("=================Test HalveDecoder====================")
# ## create and show data for test
# test_input = np.random.uniform(size=(3, 8)) # (B, F)

# def halve_fn(x, training=False):
#     halve = HalveDecoder(dim_in=8, dim_out=1, activation=jax.nn.relu)(x) # type: ignore
#     return halve

# halve_fn = jax.vmap(halve_fn, in_axes=(0,))
# ts_halve_fn = hk.transform(halve_fn, apply_rng=True)
# halve_params = ts_halve_fn.init(rng, test_input)
# halve_out = jit(ts_halve_fn.apply)(halve_params, rng, test_input)
# print(f"halve_out: shape of {halve_out.shape}, value of \n{halve_out}")

# _name = 'halve_decoder/~/mlp/~/linear'
# _params = hk.data_structures.to_mutable_dict(halve_params)
# ms_halve_fn = ms_HalveDecoder(dim_in=8, dim_out=1, activation=ms_nn.ReLU())
# ms_params = {}
# ms_params['output.mlp.0.weight'] = \
#     ms.Parameter(Tensor(np.array(_params[f'{_name}_0']['w']).T, ms.float32))
# ms_params['output.mlp.0.bias'] = \
#     ms.Parameter(Tensor(np.array(_params[f'{_name}_0']['b']), ms.float32))
# ms_params['output.mlp.1.weight'] = \
#     ms.Parameter(Tensor(np.array(_params[f'{_name}_1']['w']).T, ms.float32))
# ms_params['output.mlp.1.bias'] = \
#     ms.Parameter(Tensor(np.array(_params[f'{_name}_1']['b']), ms.float32))
# load_param_into_net(ms_halve_fn, ms_params)

# ms_test_input = Tensor(test_input, ms.float32)
# ms_halve_out = ms_halve_fn(ms_test_input)
# print(f"ms_halve_out: shape of {ms_halve_out.shape}, value of \n{ms_halve_out}")

# halve_same = np.allclose(halve_out, ms_halve_out.asnumpy(), atol=EPSILON)
# print(f"halve decoder is the same: {halve_same} in accuracy of {EPSILON}.")
# halve_diff = halve_out - ms_halve_out.asnumpy()
# halve_norm = np.linalg.norm(halve_diff)
# print(f"norm of diff: {halve_norm}")

# test residual decoder
print("=================Test ResDecoder====================")
## create and show data for test
test_input = np.random.uniform(size=(3, 8)) # (B, F)

def res_fn(x, training=False):
    res = ResDecoder(dim_in=8, dim_out=1, activation=jax.nn.relu, n_layers=1)(x) # type: ignore
    return res

res_fn = jax.vmap(res_fn, in_axes=(0,))
ts_res_fn = hk.transform(res_fn, apply_rng=True)
res_params = ts_res_fn.init(rng, test_input)
res_out = jit(ts_res_fn.apply)(res_params, rng, test_input)
print(f"res_out: shape of {res_out.shape}, value of \n{res_out}")

_params = hk.data_structures.to_mutable_dict(res_params)
ms_res_fn = ms_ResDecoder(dim_in=8, dim_out=1, activation=ms_nn.ReLU(), n_layers=1)
ms_params = {}
ms_params["output.0.preact_dense1.dense.weight"] = \
    ms.Parameter(Tensor(np.array(_params['residual_decoder/~/residual_decoder_0/~/linear_A']['w']).T, ms.float32))
ms_params["output.0.preact_dense1.dense.bias"] = \
    ms.Parameter(Tensor(np.array(_params['residual_decoder/~/residual_decoder_0/~/linear_A']['b']), ms.float32))
ms_params["output.0.preact_dense2.dense.weight"] = \
    ms.Parameter(Tensor(np.array(_params['residual_decoder/~/residual_decoder_0/~/linear_B']['w']).T, ms.float32))
ms_params["output.0.preact_dense2.dense.bias"] = \
    ms.Parameter(Tensor(np.array(_params['residual_decoder/~/residual_decoder_0/~/linear_B']['b']), ms.float32))
ms_params["output.1.dense.weight"] = \
    ms.Parameter(Tensor(np.array(_params['residual_decoder/~/last_linear']['w']).T, ms.float32))
ms_params["output.1.dense.bias"] = \
    ms.Parameter(Tensor(np.array(_params['residual_decoder/~/last_linear']['b']), ms.float32))

load_param_into_net(ms_res_fn, ms_params)

ms_test_input = Tensor(test_input, ms.float32)
ms_res_out = ms_res_fn(ms_test_input)
print(f"ms_res_out: shape of {ms_res_out.shape}, value of \n{ms_res_out}")

res_same = np.allclose(res_out, ms_res_out.asnumpy(), atol=EPSILON)
print(f"res decoder is the same: {res_same} in accuracy of {EPSILON}.")
res_diff = res_out - ms_res_out.asnumpy()
res_norm = np.linalg.norm(res_diff)
print(f"norm of diff: {res_norm}")